# -*- coding: utf-8 -*-
# @File     : detect_main.py
# @Author   : bingjia
# @Time     : 2020/9/22 16:43
# @Desc     : 检测调用接口

from django.apps import apps
from django.conf import settings

from .utils.utils import *
from .utils.datasets import *
from .models.experimental import *
from .cfg.detect_cfg import defect_cfg
from .utils.generate_pdf import PDFGenerator
from .utils.get_image_attr import get_image_attr
from .utils.change_str_datetime import change_str_datetime
from .utils.get_project_report_data import get_project_report_data

import torch.backends.cudnn as cudnn


def detect_main(input_path, output_path, project_id):
    """
    检测函数入口
    :return:
    """
    "配置源文件路径和输出路径"
    defect_cfg.source = input_path
    defect_cfg.output = output_path

    project_model = apps.get_model("projects", "Project").objects.get(id=project_id)
    report_path = os.path.join(settings.ROOT_PATH, settings.STATIC_URL[1:-1], settings.PROJECT_PATH,
                               str(project_id), "file")

    project_ref = project_model.slug

    if not os.path.exists(report_path):
        os.makedirs(report_path)

    "更新项目状态"
    apps.get_model("projects", "Project").objects.filter(id=project_id).update(analysis_status="1")

    with torch.no_grad():
        detect(project_ref)

    "获取报告数据"
    data = get_project_report_data(project_ref=project_ref)

    "生成pdf报告"
    report_name = project_model.name
    gen_pdf = PDFGenerator(os.path.join(report_path, "{}.pdf".format(report_name)), data)
    gen_pdf.gen_pv_report_pdf()

    "更新项目状态"
    apps.get_model("projects", "Project").objects.filter(id=project_id) \
        .update(analysis_status="2", view_status="1",
                report_path="/static/project/{}/file/{}.pdf".format(project_id, report_name))


def detect(project_ref):
    """
    检测函数
    :return:
    """

    "数据库"
    defectModel = apps.get_model("checks", "Defect")
    projectModel = apps.get_model("projects", "Project")
    modelInfoModel = apps.get_model("trains", "ModelInfo")
    defectTypeModel = apps.get_model("checks", "DefectType")

    "初始化参数"
    out, source, weights, view_img, save_txt, imgsz = \
        defect_cfg.output, defect_cfg.source, defect_cfg.weights, defect_cfg.view_img, defect_cfg.save_txt, defect_cfg.img_size
    webcam = source == '0' or source.startswith('rtsp') or source.startswith('http') or source.endswith('.txt')

    device = torch_utils.select_device(defect_cfg.device)
    if os.path.exists(out):
        shutil.rmtree(out)  # delete output folder
    os.makedirs(out)  # make new output folder
    half = device.type != 'cpu'

    "加载模型"
    model = attempt_load(weights, map_location=device)  # load FP32 model
    imgsz = check_img_size(imgsz, s=model.stride.max())  # check img_size
    if half:
        model.half()  # to FP16

    "分类器"
    classify = False
    if classify:
        modelc = torch_utils.load_classifier(name='resnet101', n=2)  # initialize
        modelc.load_state_dict(torch.load('weights/resnet101.pt', map_location=device)['model'])  # load weights
        modelc.to(device).eval()

    "设置dataloader"
    vid_path, vid_writer = None, None
    if webcam:
        view_img = True
        cudnn.benchmark = True  # set True to speed up constant image size inference
        dataset = LoadStreams(source, img_size=imgsz)
    else:
        save_img = True
        dataset = LoadImages(source, img_size=imgsz)

    "设置检测名字和边框颜色"
    names = model.module.names if hasattr(model, 'module') else model.names
    colors = [[random.randint(0, 255) for _ in range(3)] for _ in range(len(names))]

    "开始推理"
    t0 = time.time()
    img = torch.zeros((1, 3, imgsz, imgsz), device=device)  # init img
    _ = model(img.half() if half else img) if device.type != 'cpu' else None  # run once

    current = 0
    project = projectModel.objects.get(slug=project_ref)

    for path, img, im0s, vid_cap in dataset:
        img = torch.from_numpy(img).to(device)
        img = img.half() if half else img.float()  # uint8 to fp16/32
        img /= 255.0  # 0 - 255 to 0.0 - 1.0
        if img.ndimension() == 3:
            img = img.unsqueeze(0)

        "推理"
        t1 = torch_utils.time_synchronized()
        pred = model(img, augment=defect_cfg.augment)[0]

        "使用NMS"
        pred = non_max_suppression(pred, defect_cfg.conf_thres, defect_cfg.iou_thres, classes=defect_cfg.classes,
                                   agnostic=defect_cfg.agnostic_nms)
        t2 = torch_utils.time_synchronized()

        "分类"
        if classify:
            pred = apply_classifier(pred, modelc, img, im0s)

        # Process detections
        for i, det in enumerate(pred):  # detections per image
            if webcam:  # batch_size >= 1
                p, s, im0 = path[i], '%g: ' % i, im0s[i].copy()
            else:
                p, s, im0 = path, '', im0s

            save_path = str(Path(out) / Path(p).name)
            txt_path = str(Path(out) / Path(p).stem) + ('_%g' % dataset.frame if dataset.mode == 'video' else '')
            s += '%gx%g ' % img.shape[2:]  # print string
            gn = torch.tensor(im0.shape)[[1, 0, 1, 0]]  # normalization gain whwh
            if det is not None and len(det):
                # Rescale boxes from img_size to im0 size
                det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()

                # Print results
                for c in det[:, -1].unique():
                    n = (det[:, -1] == c).sum()  # detections per class
                    s += '%g %ss, ' % (n, names[int(c)])  # add to string

                "写入result"
                for *xyxy, conf, cls in det:
                    if save_txt:  # Write to file
                        xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()  # normalized xywh
                        with open(txt_path + '.txt', 'a') as f:
                            f.write(('%g ' * 5 + '\n') % (cls, *xywh))  # label format

                    if save_img or view_img:  # Add bbox to image
                        label = '%s %.2f' % (names[int(cls)], conf)
                        plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=1)

                    "获取图片属性信息"
                    shooting_time, lng, lat = get_image_attr(p)
                    "数据库数据初始化"
                    defect_name = names[int(cls)]
                    _, image_name = os.path.split(p)
                    model_info = modelInfoModel.objects.get(network_model__id=1, image_type="0")

                    defect_type = defectTypeModel.objects.get(name=defect_name, type="0")

                    "将缺陷数据保存到数据库"
                    defectModel.objects.create(
                        model_info=model_info,
                        defect=defect_type,
                        project=project,
                        location=str(lat) + ',' + str(lng),
                        shoot_date=change_str_datetime(shooting_time),
                        defect_image=image_name,
                        thermal=image_name,
                        visible=image_name,
                    )

            "填充分析进度"
            current = current + 1
            project.analysis_progress = int(float('%.2f' % (current / len(dataset))) * 100)
            project.save()
            # Print time (inference + NMS)
            print('%sDone. (%.3fs)' % (s, t2 - t1))

            # Stream results
            if view_img:
                cv2.imshow(p, im0)
                if cv2.waitKey(1) == ord('q'):  # q to quit
                    raise StopIteration

            # Save results (image with detections)
            if save_img:
                if dataset.mode == 'images':
                    cv2.imwrite(save_path, im0)
                else:
                    if vid_path != save_path:  # new video
                        vid_path = save_path
                        if isinstance(vid_writer, cv2.VideoWriter):
                            vid_writer.release()  # release previous video writer

                        fourcc = 'mp4v'  # output video codec
                        fps = vid_cap.get(cv2.CAP_PROP_FPS)
                        w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
                        h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
                        vid_writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*fourcc), fps, (w, h))
                    vid_writer.write(im0)

    if save_txt or save_img:
        print('Results saved to %s' % out)
        if platform == 'darwin' and not defect_cfg.update:  # MacOS
            os.system('open ' + save_path)

    print('Done. (%.3fs)' % (time.time() - t0))
